package com.pearl.security.multiloginreject.demo.session.memery;

import cn.hutool.core.util.StrUtil;
import com.pearl.security.multiloginreject.demo.util.LoginUtils;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.context.ApplicationListener;
import org.springframework.core.log.LogMessage;
import org.springframework.security.core.session.SessionDestroyedEvent;
import org.springframework.security.core.session.SessionInformation;
import org.springframework.security.core.session.SessionRegistry;
import org.springframework.util.Assert;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.CopyOnWriteArraySet;

/**
 * 在实现方案中，
 * @author TD
 * @version 1.0
 * @date 2024/5/15
 */
public class PearlSessionRegistryImpl implements SessionRegistry, ApplicationListener<SessionDestroyedEvent> {
    protected final Log logger = LogFactory.getLog(PearlSessionRegistryImpl.class);

    /**
     * 存放当前登录用户标识 和 会话ID集合
     * PC_1235456-》4E518646B47B6584C3BD493C38C9F54B，F848E8ABA968E1D7BF129A71FE9A4D48
     */
    private final ConcurrentMap<String, Set<String>> principals;

    /**
     * 存放会话ID 和 会话
     * 4E518646B47B6584C3BD493C38C9F54B-》SessionInformation
     */
    private final Map<String, SessionInformation> sessionIds;


    public PearlSessionRegistryImpl() {
        this.principals = new ConcurrentHashMap<>();
        this.sessionIds = new ConcurrentHashMap<>();
    }

    public PearlSessionRegistryImpl(ConcurrentMap<String, Set<String>> principals, Map<String, SessionInformation> sessionIds) {
        this.principals = principals;
        this.sessionIds = sessionIds;
    }

    public List<Object> getAllPrincipals() {
        return new ArrayList<>(this.principals.keySet());
    }

    /**
     * 获取当前登录用户的标识 eg:com.pearl.security.multiloginreject.demo.session.concurrent.PearlSessionRegistryImpl_PC_123456
     *
     * @param principal com.pearl.security.multiloginreject.demo.content.PearlUserDetails
     * @return key
     */
    public String getKey(Object principal) {
        String key = PearlSessionRegistryImpl.class.getName() + StrUtil.UNDERLINE;
        key += LoginUtils.parseDevice(principal) + StrUtil.UNDERLINE + LoginUtils.parseUserId(principal);
        return key;
    }

    /**
     * 查询所有会话
     *
     * @param principal              org.springframework.security.core.userdetails.User
     * @param includeExpiredSessions 是否包含失效的会话
     * @return 所有会话
     */
    public List<SessionInformation> getAllSessions(Object principal, boolean includeExpiredSessions) {
        // 查询当前用户的所有会话ID
        Set<String> sessionsUsedByPrincipal = this.principals.get(getKey(principal));
        if (sessionsUsedByPrincipal == null) {
            return Collections.emptyList();
        }
        // 根据会话ID查询会话
        List<SessionInformation> list = new ArrayList<>(sessionsUsedByPrincipal.size());
        for (String sessionId : sessionsUsedByPrincipal) {
            SessionInformation sessionInformation = getSessionInformation(sessionId);
            if (sessionInformation == null) {
                continue;
            }
            if (includeExpiredSessions || !sessionInformation.isExpired()) {
                list.add(sessionInformation);
            }
        }
        return list;
    }

    /**
     * 根据会话ID查询会话
     *
     * @param sessionId 会话ID
     * @return 会话
     */
    public SessionInformation getSessionInformation(String sessionId) {
        Assert.hasText(sessionId, "SessionId required as per interface contract");
        return this.sessionIds.get(sessionId);
    }

    /**
     * 会话销毁
     *
     * @param event 会话销毁时间
     */
    public void onApplicationEvent(SessionDestroyedEvent event) {
        String sessionId = event.getId();
        this.removeSessionInformation(sessionId);
    }

    /**
     * 每次访问时，刷新最后访问时间
     *
     * @param sessionId 会话ID
     */
    @Override
    public void refreshLastRequest(String sessionId) {
        Assert.hasText(sessionId, "SessionId required as per interface contract");
        SessionInformation info = getSessionInformation(sessionId);
        if (this.logger.isDebugEnabled()) {
            this.logger.debug(LogMessage.format("Refresh Last Request session %s", sessionId));
        }
        if (info != null) {
            info.refreshLastRequest();
        }
    }

    /**
     * 注册新会话
     *
     * @param sessionId 会话ID
     * @param principal 登录用户（org.springframework.security.core.userdetails.User）
     */
    @Override
    public void registerNewSession(String sessionId, Object principal) {
        Assert.hasText(sessionId, "SessionId required as per interface contract");
        Assert.notNull(principal, "Principal required as per interface contract");
        if (getSessionInformation(sessionId) != null) {
            removeSessionInformation(sessionId);
        }
        if (this.logger.isDebugEnabled()) {
            this.logger.debug(LogMessage.format("Registering session %s, for principal %s", sessionId, principal));
        }
        this.sessionIds.put(sessionId, new SessionInformation(principal, sessionId, new Date()));
        this.principals.compute(getKey(principal), (key, sessionsUsedByPrincipal) -> {
            if (sessionsUsedByPrincipal == null) {
                sessionsUsedByPrincipal = new CopyOnWriteArraySet<>();
            }
            sessionsUsedByPrincipal.add(sessionId);
            this.logger.trace(LogMessage.format("Sessions used by '%s' : %s", principal, sessionsUsedByPrincipal));
            return sessionsUsedByPrincipal;
        });
    }

    /**
     * 移除会话
     *
     * @param sessionId 会话ID
     */
    @Override
    public void removeSessionInformation(String sessionId) {
        Assert.hasText(sessionId, "SessionId required as per interface contract");
        SessionInformation info = getSessionInformation(sessionId);
        if (info == null) {
            return;
        }
        if (this.logger.isTraceEnabled()) {
            this.logger.debug("Removing session " + sessionId + " from set of registered sessions");
        }
        this.sessionIds.remove(sessionId);
        this.principals.computeIfPresent(getKey(info.getPrincipal()), (key, sessionsUsedByPrincipal) -> {
            this.logger.debug(
                    LogMessage.format("Removing session %s from principal's set of registered sessions", sessionId));
            sessionsUsedByPrincipal.remove(sessionId);
            if (sessionsUsedByPrincipal.isEmpty()) {
                // No need to keep object in principals Map anymore
                this.logger.debug(LogMessage.format("Removing principal %s from registry", info.getPrincipal()));
                sessionsUsedByPrincipal = null;
            }
            this.logger.trace(
                    LogMessage.format("Sessions used by '%s' : %s", info.getPrincipal(), sessionsUsedByPrincipal));
            return sessionsUsedByPrincipal;
        });
    }
}
